function [residual,at,ct,ht,inct,zt,labinct,capinct,totalinct,cpnext] = to_aiyagari_stationary(r0,crit,I,T,Amat,Zmat,alpha,b,delta,rhovector,sharesuper,varphi,A0,C0,H)
% this function
% (1) solves for the consumption decision rule via policy function iteration, given an
% intereste rate, r0
% (2) simulates the stationary equilibrium associated with the interest
% rate, r0
% (3) (i) returns the residual between the given interes rate r0 and the one
% implied by the stationary aggregate capital and labor supply.
% (ii) returns as an optional output the wealth distribution.
% get dimensions of the grid
[M,N] = size(Amat);

% get productivity realizations from the first row
z1 = Zmat(1,1);
z2 = Zmat(1,2);
z3 = Zmat(1,3);
z4 = Zmat(1,4);
z5 = Zmat(1,5);

% re-create the Markov-chain transition probability matrix
rho11 = rhovector(1,1);
rho12 = rhovector(1,2);
rho13 = rhovector(1,3);
rho14 = rhovector(1,4);
rho15 = rhovector(1,5);
rho21 = rhovector(1,6);
rho22 = rhovector(1,7);
rho23 = rhovector(1,8);
rho24 = rhovector(1,9);
rho25 = rhovector(1,10);
rho31 = rhovector(1,11);
rho32 = rhovector(1,12);
rho33 = rhovector(1,13);
rho34 = rhovector(1,14);
rho35 = rhovector(1,15);
rho41 = rhovector(1,16);
rho42 = rhovector(1,17);
rho43 = rhovector(1,18);
rho44 = rhovector(1,19);
rho45 = rhovector(1,20);
rho51 = rhovector(1,21);
rho52 = rhovector(1,22);
rho53 = rhovector(1,23);
rho54 = rhovector(1,24);
rho55 = rhovector(1,25);

sharenotsuper = 1-sharesuper;

threshold11=rho11;
threshold12=rho11+rho12;
threshold13=rho11+rho12+rho13;
threshold14=rho11+rho12+rho13+rho14;

threshold21=rho21;
threshold22=rho21+rho22;
threshold23=rho21+rho22+rho23;
threshold24=rho21+rho22+rho23+rho24;

threshold31=rho31;
threshold32=rho31+rho32;
threshold33=rho31+rho32+rho33;
threshold34=rho31+rho32+rho33+rho34;

threshold41=rho41;
threshold42=rho41+rho42;
threshold43=rho41+rho42+rho43;
threshold44=rho41+rho42+rho43+rho44;

threshold51=rho51;
threshold52=rho51+rho52;
threshold53=rho51+rho52+rho53;
threshold54=rho51+rho52+rho53+rho54;

threshold= [threshold11 threshold12 threshold13 threshold14; threshold21 threshold22 threshold23 threshold24; threshold31 threshold32 threshold33 threshold34; threshold41 threshold42 threshold43 threshold44];


% compute the wage according to marginal pricing and a Cobb-Douglas 
% production function
w0 = (1-alpha)*(alpha/(r0+delta))^(alpha/(1-alpha));

% initial guess on future consumption (consume asset income plus labor
% income from working h=1.
cp0 = r0*Amat+Zmat*w0;

%%% iterate on the consumption decision rule

% distance between two successive functions
% start with a distance above the convergence criterion
dist    = crit+1;
% maximum iterations (avoids infinite loops)
maxiter = 10^(5);
% counter
iter    = 1;

fprintf('Inner loop on policy function iteration is initiated ... \n');
    
while (dist>crit&&iter<maxiter)

% derive current consumption via C0(.) function
c0 = C0(cp0,r0);

% derive current assets
a0 = A0(Amat,Zmat,c0,r0,w0);

%%% update the guess for the consumption decision

% consumption decision rule for a binding borrowing constraint
% can be solved as a quadratic equation
% c^(2)-((1+r)a+b)c-(yw)^(1+2/3) = 0
% general formula and notation: ax^(2)+bx+c = 0
% x = (-b+sqrt(b^2-4ac))/2a
cpbind = ((1+r0)*Amat+b)/2+sqrt(((1+r0)*Amat+b).^(2)+4*(Zmat*w0).^(1+varphi))/2;
% % slow alternative: rootfinding gives the same result
% cpbind  = zeros(M,N);
% options = optimset('Display','off');
% cpbind(:,1) = fsolve(@(c) (1+r0)*Amat(:,1)+H(c,Y(1),w0).*Z(1).*w0-c+b,cp0(:,1),options);
% cpbind(:,2) = fsolve(@(c) (1+r0)*Amat(:,2)+H(c,Y(2),w0).*Z(2).*w0-c+b,cp0(:,2),options);
% cpbind(:,3) = fsolve(@(c) (1+r0)*Amat(:,3)+H(c,Y(3),w0).*Z(3).*w0-c+b,cp0(:,3),options);
% cpbind(:,4) = fsolve(@(c) (1+r0)*Amat(:,4)+H(c,Y(4),w0).*Z(4).*w0-c+b,cp0(:,4),options);
% cpbind(:,5) = fsolve(@(c) (1+r0)*Amat(:,5)+H(c,Y(5),w0).*Z(5).*w0-c+b,cp0(:,5),options);



% consumption for nonbinding borrowing constraint
cpnon = zeros(M,N);
% interpolation conditional on productivity realization
% instead of extrapolation use the highest value in a0
cpnon(:,1)  = interp1(a0(:,1),c0(:,1),Amat(:,1),'spline');
cpnon(:,2)  = interp1(a0(:,2),c0(:,2),Amat(:,2),'spline');
cpnon(:,3)  = interp1(a0(:,3),c0(:,3),Amat(:,3),'spline');
cpnon(:,4)  = interp1(a0(:,4),c0(:,4),Amat(:,4),'spline');
cpnon(:,5)  = interp1(a0(:,5),c0(:,5),Amat(:,5),'spline');


% merge the two, separate grid points that induce a binding borrowing constraint
% for the future asset level (the first observation of the endogenous current asset grid is the 
% threshold where the borrowing constraint starts binding, for all lower values it will also bind
% as the future asset holdings are monotonically increasing in the current
% asset holdings).
cpnext(:,1) = (Amat(:,1)>a0(1,1)).*cpnon(:,1)+(Amat(:,1)<=a0(1,1)).*cpbind(:,1);
cpnext(:,2) = (Amat(:,2)>a0(1,2)).*cpnon(:,2)+(Amat(:,2)<=a0(1,2)).*cpbind(:,2);
cpnext(:,3) = (Amat(:,3)>a0(1,3)).*cpnon(:,3)+(Amat(:,3)<=a0(1,3)).*cpbind(:,3);
cpnext(:,4) = (Amat(:,4)>a0(1,4)).*cpnon(:,4)+(Amat(:,4)<=a0(1,4)).*cpbind(:,4);
cpnext(:,5) = (Amat(:,5)>a0(1,5)).*cpnon(:,5)+(Amat(:,5)<=a0(1,5)).*cpbind(:,5);



% distance measure
dist = norm((cpnext-cp0)./cp0);

% display every 100th iteration
if mod(iter,100) == 1
    fprintf('Inner loop on policy function iteration, iteration: %3i, Norm: %2.6f \n',[iter,dist]);
end

% increase counter
iter = iter+1;

% update the guess on the consumption function
cp0 = cpnext;


end

fprintf('Inner loop on policy function iteration is done. \n');

fprintf('Starting simulation... \n');

%%% simulate the stationary wealth distribution

% initialize variables
at              = zeros(I,T+1);
zt              = zeros(I,T+1);
ct              = zeros(I,T+1);
ht              = zeros(I,T+1);
inct            = zeros(I,T+1);
capinct         = zeros(I,T+1);
labinct         = zeros(I,T+1);
totalinct       = zeros(I,T+1);
altlaborinct    = zeros(I,T+1);
% period T+1 is to be discarded

at(:,1) = 1;            % initial asset level

for t=1:T
   % draw uniform random numbers across individuals
   s1 = unifrnd(0,1,I,1);       %shocks for individuals in state y1
   s2 = unifrnd(0,1,I,1);       %shocks for individuals in state y2
   s3 = unifrnd(0,1,I,1);       %shocks for individuals in state y3
   s4 = unifrnd(0,1,I,1);       %shocks for individuals in state y4
   s5 = unifrnd(0,1,I,1);       %shocks for individuals in state y5
 
   if t>1   % transition for t>1
       zt(:,t) = ((s1<=threshold11)&(zt(:,t-1)==z1)).*z1+ (s1>threshold11 & s1<=threshold12& zt(:,t-1)==z1).*z2+ (s1>threshold12 & s1<=threshold13& zt(:,t-1)==z1).*z3+ (s1>threshold13 & s1<=threshold14& zt(:,t-1)==z1).*z4+ (s1>threshold14 & zt(:,t-1)==z1).*z5...
                +((s2<=threshold21)&(zt(:,t-1)==z2)).*z1+ (s2>threshold21 & s2<=threshold22& zt(:,t-1)==z2).*z2+ (s2>threshold22 & s2<=threshold23& zt(:,t-1)==z2).*z3+ (s2>threshold23 & s2<=threshold24& zt(:,t-1)==z2).*z4+ (s2>threshold24 & zt(:,t-1)==z2).*z5...    
                +((s3<=threshold31)&(zt(:,t-1)==z3)).*z1+ (s3>threshold31 & s3<=threshold32& zt(:,t-1)==z3).*z2+ (s3>threshold32 & s3<=threshold33& zt(:,t-1)==z3).*z3+ (s3>threshold33 & s3<=threshold34& zt(:,t-1)==z3).*z4+ (s3>threshold34 & zt(:,t-1)==z3).*z5...    
                +((s4<=threshold41)&(zt(:,t-1)==z4)).*z1+ (s4>threshold41 & s4<=threshold42& zt(:,t-1)==z4).*z2+ (s4>threshold42 & s4<=threshold43& zt(:,t-1)==z4).*z3+ (s4>threshold43 & s4<=threshold44& zt(:,t-1)==z4).*z4+ (s4>threshold44 & zt(:,t-1)==z4).*z5...    
                +((s5<=threshold51)&(zt(:,t-1)==z5)).*z1+ (s5>threshold51 & s5<=threshold52& zt(:,t-1)==z5).*z2+ (s5>threshold52 & s5<=threshold53& zt(:,t-1)==z5).*z3+ (s5>threshold53 & s5<=threshold54& zt(:,t-1)==z5).*z4+ (s5>threshold54 & zt(:,t-1)==z5).*z5;    

            
   else     % random allocation in t=0
       zt(:,t) = (s2<sharenotsuper & s1<=1/3).*z1+(s2<sharenotsuper & (s1>1/3)&(s1<2/3)).*z2+(s2<sharenotsuper & s1>=2/3).*z3 + (s2>=sharenotsuper & s3<=1/2).*z4+ (s2>=sharenotsuper & s3>1/2).*z5 ; 
   end

   % consumption
   ct(:,t)   = (zt(:,t)==z1).*interp1(Amat(:,1),cp0(:,1),at(:,t),'spline')+(zt(:,t)==z2).*interp1(Amat(:,2),cp0(:,2),at(:,t),'spline')+(zt(:,t)==z3).*interp1(Amat(:,3),cp0(:,3),at(:,t),'spline')+(zt(:,t)==z4).*interp1(Amat(:,4),cp0(:,4),at(:,t),'spline')+(zt(:,t)==z5).*interp1(Amat(:,5),cp0(:,5),at(:,t),'spline');

   % labor supply
   ht(:,t)   = H(ct(:,t),zt(:,t),w0);

   % future assets
   at(:,t+1) = (1+r0)*at(:,t)+ht(:,t).*zt(:,t).*w0-ct(:,t); 
   
   % income
   capinct(:,t+1)=r0.*at(:,t);
   labinct(:,t)=w0.*ht(:,t).*zt(:,t);
   % alternatively labor income can be as below
   % altlaborinct(:,t)= at(:,t+1)+ct(:,t)-(1+r0)*at(:,t);
   totalinct(:,t)=labinct(:,t)+capinct(:,t);
end

fprintf('simulation done... \n');

% compute aggregates from market clearing
% K = mean(mean(at(:,T-100:T)));
% L = mean(mean(yt(:,T-100:T).*ht(:,T-100:T)));
% H = mean(mean(ht(:,T-100:T)));
% INC = mean(mean(inct(:,T-100:T)));
% C = mean(mean(ct(:,T-100:T)));
% LI= mean(mean(labinct(:,T-100:T)));
% CI= mean(mean(capinct(:,T-100:T)));
% TI= mean(mean(totalinct(:,T-100:T)));
% r = alpha*(K/L)^(alpha-1)-delta;
% w = (1-alpha)*(alpha/(r+delta))^(alpha/(1-alpha));


K   = mean(mean(at(:,T-100:T)));
L   = mean(mean(zt(:,T-100:T).*ht(:,T-100:T)));
H   = mean(mean(ht(:,T-100:T)));
C   = mean(mean(ct(:,T-100:T)));
LI  = mean(mean(labinct(:,T-100:T)));
CI  = mean(mean(capinct(:,T-100:T)));
TI  = mean(mean(totalinct(:,T-100:T)));
%ALTLI = mean(altlaborinct(:,T-100:T));
r   = alpha*(K/L)^(alpha-1)-delta;
w   = (1-alpha)*(alpha/(r+delta))^(alpha/(1-alpha));



fprintf('average physical capital is: %2.3f \n',K)
fprintf('average effective labor is: %2.3f \n',L)
fprintf('average labor is: %2.3f \n',H)
fprintf('average consumption is: %2.3f \n',C)
fprintf('real interest rate is: %2.3f \n',r)
fprintf('wage rate is: %2.3f \n',w)
fprintf('average capital income is: %2.3f \n',CI)
fprintf('average labor income is: %2.3f \n',LI)
fprintf('average total income is: %2.3f \n',TI)
% compute the distance between the two
residual = (r-r0)/r0;

end
